package com.btl.service.ai.support.rag.generation.assistant;

import com.btl.service.ai.support.rag.generation.GenerationModelEnum;
import com.btl.service.ai.support.rag.store.rdb.mapper.ChatHistoryMapper;
import com.btl.service.ai.support.rag.store.rdb.table.ChatHistory;
import com.fasterxml.jackson.core.type.TypeReference;
import dev.langchain4j.data.message.ChatMessageType;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.TokenCountEstimator;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.Result;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.UUID;

@Slf4j
@Service
public abstract class GenerationAssistantStrategy {


    @Resource
    private ChatHistoryMapper chatHistoryMapper;

    /**
     * 生成指定类型的结果
     *
     * @param memoryId      记忆id
     * @param systemMessage 系统消息
     * @param userMessage   用户消息
     * @param query         查询
     * @param responseType  响应类型
     * @return 结果
     */
    public abstract <T> Result<T> generateJson(
            Long memoryId,
            String systemMessage,
            String userMessage,
            Query query,
            Class<T> responseType);

    /**
     * 生成指定类型的结果
     *
     * @param memoryId      记忆id
     * @param systemMessage 系统消息
     * @param userMessage   用户消息
     * @param query         查询
     * @param typeRef       响应类型
     * @return 结果
     */
    public abstract <F> Result<F> generateJson(
            Long memoryId,
            String systemMessage,
            String userMessage,
            Query query,
            TypeReference<F> typeRef);

    /**
     * 生成字符串类型结果
     */
    public abstract Result<String> generate(
            Long memoryId,
            String systemMessage,
            String userMessage,
            Query query);

    /**
     * 模型
     */
    protected abstract GenerationModelEnum model();

    /**
     * 构建 AiService
     *
     * @param chatModel                      模型
     * @param redisChatMemoryStore           缓存
     * @param huggingFaceTokenCountEstimator 模型
     * @return AIService
     */
    protected AssistantService buildAiService(ChatModel chatModel, ChatMemoryStore redisChatMemoryStore, TokenCountEstimator huggingFaceTokenCountEstimator) {
        return AiServices
                .builder(AssistantService.class)
                .chatModel(chatModel)
                .chatMemoryProvider(mid ->
                        TokenWindowChatMemory
                                .builder()
                                .id(mid)
                                .chatMemoryStore(redisChatMemoryStore)
                                .maxTokens(10000, huggingFaceTokenCountEstimator)
                                .build()
                )
                .chatRequestTransformer(chatRequest -> {
                    return chatRequest.toBuilder()
                            .build();
                })
                .build();
    }


    protected Result<String> generate(AssistantService assistantService, Long memoryId, String systemMessage, String userMessage, Query query) {
        String aqId = UUID.randomUUID()
                .toString();
        ChatHistory chatHistoryRequest = new ChatHistory();
        chatHistoryRequest.setMemoryId(memoryId);
        chatHistoryRequest.setContent(query.text());
        chatHistoryRequest.setRole(ChatMessageType.USER.name());
        chatHistoryRequest.setAqId(aqId);
        chatHistoryMapper.insert(chatHistoryRequest);
        Result<String> answer = assistantService
                .answer(memoryId, systemMessage, userMessage, query.text());
        ChatHistory chatHistoryResponse = new ChatHistory();
        chatHistoryResponse.setMemoryId(memoryId);
        chatHistoryResponse.setContent(answer.content());
        chatHistoryResponse.setRole(ChatMessageType.AI.name());
        chatHistoryResponse.setAqId(aqId);
        chatHistoryMapper.insert(chatHistoryResponse);
        return answer;
    }

}
